#!/usr/bin/env python
# -*- coding: utf-8 -*-

import re
import numpy as np
import matplotlib.pyplot as plt

def parse_log_line(log_line):
    """
    解析日志行，检查是否包含 'training' 或 'validation'，并提取 Loss、SSIM 和 PSNR 值。

    :param log_line: 日志行字符串
    :return: 包含 'training' 或 'validation' 的字符串，以及提取的 Loss、SSIM 和 PSNR 值
    """
    # 定义正则表达式模式
    training_pattern = r'training: Epoch \[\d+\], Start \[\d+\], Step \[\d+/\d+\], Loss: ([\d.]+),Time:\[[\d.]+\], SSIM: ([\d.]+), PSNR: ([\d.]+)'
    validation_pattern = r'validation: Epoch \[\d+\],\s*Validation Loss: ([\d.]+), SSIM: ([\d.]+), PSNR: ([\d.]+)'
    
    # 搜索匹配的日志行
    training_match = re.search(training_pattern, log_line)
    validation_match = re.search(validation_pattern, log_line)
    
    if training_match:
        log_type = 'training'
        loss = float(training_match.group(1))
        ssim = float(training_match.group(2))
        psnr = float(training_match.group(3))
    elif validation_match:
        log_type = 'validation'
        loss = float(validation_match.group(1))
        ssim = float(validation_match.group(2))
        psnr = float(validation_match.group(3))
    else:
        log_type = None
        loss = None
        ssim = None
        psnr = None
    
    return log_type, loss, ssim, psnr


def read_and_parse_log_file(log_path):
    """
    读取日志文件并解析每一行，将训练和验证的 Loss、SSIM 和 PSNR 分别存储在不同的数组中。

    :param log_path: 日志文件路径
    :return: 训练和验证的 Loss、SSIM 和 PSNR 的 NumPy 数组
    """
    training_loss = []
    training_ssim = []
    training_psnr = []
    
    validation_loss = []
    validation_ssim = []
    validation_psnr = []
    
    with open(log_path, 'r', encoding='utf-8') as file:
        for line in file:
            log_type, loss, ssim, psnr = parse_log_line(line)
            if log_type == 'training':
                training_loss.append(loss)
                training_ssim.append(ssim)
                training_psnr.append(psnr)
            elif log_type == 'validation':
                validation_loss.append(loss)
                validation_ssim.append(ssim)
                validation_psnr.append(psnr)
            else:
                print('invalid log line!')
    
    # 将列表转换为 NumPy 数组
    training_loss = np.array(training_loss)
    training_ssim = np.array(training_ssim)
    training_psnr = np.array(training_psnr)
    
    validation_loss = np.array(validation_loss)
    validation_ssim = np.array(validation_ssim)
    validation_psnr = np.array(validation_psnr)
    
    return training_loss, training_ssim, training_psnr, validation_loss, validation_ssim, validation_psnr

def plot_arrays(training_loss, training_ssim, training_psnr, validation_loss, validation_ssim, validation_psnr):
    """
    使用 matplotlib 绘制训练和验证的 Loss、SSIM 和 PSNR 数组。

    :param training_loss, training_ssim, training_psnr: 训练数据的数组
    :param validation_loss, validation_ssim, validation_psnr: 验证数据的数组
    """
    epochs = np.arange(len(training_loss)) + 1
    val_epochs = np.arange(len(validation_loss)) + 1

    fig, axes = plt.subplots(3, 2, figsize=(15, 15))

    # Training Loss
    axes[0, 0].plot(epochs, training_loss, label='Training Loss')
    axes[0, 0].set_title('Training Loss')
    axes[0, 0].set_xlabel('Epoch')
    axes[0, 0].set_ylabel('Loss')
    axes[0, 0].legend()

    # Validation Loss
    axes[0, 1].plot(val_epochs, validation_loss, label='Validation Loss', color='orange')
    axes[0, 1].set_title('Validation Loss')
    axes[0, 1].set_xlabel('Epoch')
    axes[0, 1].set_ylabel('Loss')
    axes[0, 1].legend()

    # Training SSIM
    axes[1, 0].plot(epochs, training_ssim, label='Training SSIM')
    axes[1, 0].set_title('Training SSIM')
    axes[1, 0].set_xlabel('Epoch')
    axes[1, 0].set_ylabel('SSIM')
    axes[1, 0].legend()

    # Validation SSIM
    axes[1, 1].plot(val_epochs, validation_ssim, label='Validation SSIM', color='orange')
    axes[1, 1].set_title('Validation SSIM')
    axes[1, 1].set_xlabel('Epoch')
    axes[1, 1].set_ylabel('SSIM')
    axes[1, 1].legend()

    # Training PSNR
    axes[2, 0].plot(epochs, training_psnr, label='Training PSNR')
    axes[2, 0].set_title('Training PSNR')
    axes[2, 0].set_xlabel('Epoch')
    axes[2, 0].set_ylabel('PSNR (dB)')
    axes[2, 0].legend()

    # Validation PSNR
    axes[2, 1].plot(val_epochs, validation_psnr, label='Validation PSNR', color='orange')
    axes[2, 1].set_title('Validation PSNR')
    axes[2, 1].set_xlabel('Epoch')
    axes[2, 1].set_ylabel('PSNR (dB)')
    axes[2, 1].legend()

    plt.tight_layout()
    plt.show()

if __name__ == '__main__':

    log_path = "E:/YCS_Temp/project/thesis_dataset/dataset/PA_DATA/app_0205.log"

    # 读取并解析日志文件
    training_loss, training_ssim, training_psnr, validation_loss, validation_ssim, validation_psnr = read_and_parse_log_file(log_path)

    # 绘制图表
    plot_arrays(training_loss, training_ssim, training_psnr, validation_loss, validation_ssim, validation_psnr)


